#%%
import pickle
import numpy as pd
import os
import os.path as osp
import sys
import shutil
import h5py
import numpy as np
import pickle
import cv2
import pandas as pd
import tqdm
import multiprocessing as mp
from lilab.openlabcluster_postprocess.s1a_clipNames_inplace_parse import parse_name




clippredpkl_old_perc66 = '/home/liying_lab/chenxf/ml-project/论文图表/semisupervised分类/MM_FF_Representive_K34.clippredpkl'
predpkldata = pickle.load(open(clippredpkl_old_perc66, 'rb'))
ind_rawclip = predpkldata['ind_rawclip']
cluster_labels_sub = predpkldata['cluster_labels']
clipNames_sub = predpkldata['clipNames']
df_clipNames_sub = predpkldata['df_clipNames']
cluster_labels_sub = cluster_labels_sub - cluster_labels_sub.min() + 1 #start from 1


predpklfile_iterstart = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/olc-iter3-2024-05-27/output/semisupervise-enc-2/olc-iter1-2024-05-23_semiseq2seq_pca14.clippredpkl'
predpkldata_iterstart = pickle.load(open(predpklfile_iterstart, 'rb'))
df_clipNames_iterstart = predpkldata_iterstart['df_clipNames']

project_iternext = '/DATA/taoxianming/rat/data/Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/olc-iter4-2024-05-27'
df_clipNames_sub_iterstart = df_clipNames_iterstart.loc[ind_rawclip]
assert np.all(df_clipNames_sub_iterstart.values == df_clipNames_sub.values)

df_clipNames_iterstart['raw_index'] = np.arange(len(df_clipNames_iterstart))
df_clipNames_iterstart['cluster_labels'] = predpkldata_iterstart['cluster_labels']



def get_ind_pvalue_mirror():
    # start from 0
    nK = 34
    nK_mutual = 8
    nK_mirrorhalf = (nK-nK_mutual)//2
    ind_pvalue_mirror = np.concatenate([np.arange(nK_mutual),
                                        np.arange(nK_mirrorhalf) + nK_mirrorhalf + nK_mutual,
                                        np.arange(nK_mirrorhalf) + nK_mutual])
    assert set(ind_pvalue_mirror) == set(np.arange(nK))
    
    def fun_label_mirror(x):
        return  ind_pvalue_mirror[np.array(x)]
    return ind_pvalue_mirror, fun_label_mirror

ind_pvalue_mirror, fun_label_mirror = get_ind_pvalue_mirror()

#%%
df_sort = df_clipNames_iterstart.sort_values(by=['vnake', 'startFrame', 'isBlack']).copy()
df_sort.reset_index(inplace=True, drop=True)
ind_sort = df_sort.index.values
new_label_sort = df_sort['cluster_labels'] - 1 #start from 0
new_label_W = new_label_sort[::2]
new_label_B = new_label_sort[1::2]
new_label_B_mirror = fun_label_mirror(new_label_B)


ind_mirror = (new_label_W == new_label_B_mirror) & (new_label_W>=0)
print('mean mirror', np.mean(new_label_W == new_label_B_mirror))
indx_mirror_not = (np.where(~ind_mirror)[0][:,None] * 2 + [[0,1]]).ravel()

df_sort.loc[indx_mirror_not, 'cluster_labels'] = -1
df_clipNames_iterstartb = df_sort.sort_values(by=['raw_index'])
df_clipNames_iterstartb.reset_index(inplace=True, drop=True)
np.mean(df_clipNames_iterstartb.loc[ind_rawclip]['cluster_labels'].values == cluster_labels_sub)
df_clipNames_iterstartb.loc[ind_rawclip]['cluster_labels'] = cluster_labels_sub

assert np.all(df_clipNames_iterstart[['isBlack', 'vnake', 'startFrame']].values == df_clipNames_iterstartb[['isBlack', 'vnake', 'startFrame']].values)

#%%
cluster_labels = df_clipNames_iterstartb['cluster_labels'].values
cluster_labels[cluster_labels<0] = 0
print('new annot labels %.2f'%(np.mean(cluster_labels>0)))
np.save(osp.join(project_iternext, 'label/label.npy'), cluster_labels)

fp = h5py.File(osp.join(project_iternext, 'datasets/data.h5'), 'r+')
a = fp.pop('label')
fp.create_dataset('label', data=cluster_labels)
fp.close()

